/*
 * SPDX-License-Identifier: BSD-2-Clause
 * Copyright (c) 2019, Intel Corporation
 */
#include <tss2/tss2_tctildr.h>

#include <stdio.h>            // for NULL
#include <string.h>           // for memcmp
#include <stdlib.h>

#include "session-util.h"     // for roll_nonces, SESSION, check_response_hmac
#include "sys-util.h"         // for DefineNvIndex, TPM2B_SIZE, TSS2_RETRY_EXP
#include <tss2/tss2_common.h>      // for TSS2_RC_SUCCESS, TSS2_RC, BYTE
#include <tss2/tss2_sys.h>         // for Tss2_Sys_PolicyAuthValue, TSS2L_SYS_AUT...
#include <tss2/tss2_tpm2_types.h>  // for TPM2_RH_NULL, TPMS_AUTH_COMMAND, TPM2B_...
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

#define CHECK_RETVAL(rc, expected, prefix)                                     \
    {                                                                          \
        if ((rc) != expected) {                                                \
            printf("%sFAILED [%s:%u] %s\n", prefix, __FUNCTION__, __LINE__,    \
                   Tss2_RC_Decode(rc));                                        \
            cleanup(rc);                                                       \
        } else {                                                               \
            printf("%sPASSED [%s:%u]\n", prefix, __FUNCTION__, __LINE__);      \
        }                                                                      \
    }

// Check Passed Case
#define CHECK_PASSED(rc) CHECK_RETVAL(rc, TPM2_RC_SUCCESS, "PassedCase: ")
// Check Failed Case
#define CHECK_FAILED(rc, expected) CHECK_RETVAL(rc, expected, "FailedCase: ")
// Check if the return code indicates success, clean up if not
#define CHECK_RC(rc) CHECK_RETVAL(rc, TPM2_RC_SUCCESS, "")

#define LOGMODULE test
#include "log.h"         // for LOG_ERROR, LOG_INFO

#define TPM20_INDEX_PASSWORD_TEST       0x01500020

#define NV_DATA_SIZE 4
#define NV_DATA { 0x00, 0xff, 0x55, 0xaa }
#define SECRET_SIZE 13
#define SECRET_DATA { 's', 'h', 'a', 'r', 'e', 'd', ' ', \
                      's', 'e', 'c', 'r', 'e', 't', }

TSS2_RC
create_policy (TSS2_SYS_CONTEXT *sys_ctx,
               TPM2B_DIGEST *authPolicy)
{
    TSS2_RC rc;
    SESSION *trialPolicySession = NULL;
    TPM2B_NONCE nonceCaller = { 0, };
    TPM2B_ENCRYPTED_SECRET encryptedSalt = { 0, };
    TPMT_SYM_DEF symmetric = {
        .algorithm = TPM2_ALG_NULL,
    };

    rc = create_auth_session (&trialPolicySession,
                              TPM2_RH_NULL,
                              0,
                              TPM2_RH_NULL,
                              0,
                              &nonceCaller,
                              &encryptedSalt,
                              TPM2_SE_TRIAL,
                              &symmetric,
                              TPM2_ALG_SHA256,
                              sys_ctx);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("create_auth_session failed with rc: 0x%x", rc);
        return rc;
    }
    rc = Tss2_Sys_PolicyAuthValue (sys_ctx,
                                   trialPolicySession->sessionHandle,
                                   0,
                                   0);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_PolicyAuthValue failed with rc: 0x%x", rc);
        return rc;
    }
    rc = Tss2_Sys_PolicyGetDigest (sys_ctx,
                                   trialPolicySession->sessionHandle,
                                   0,
                                   authPolicy,
                                   0);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_PolicyGetDigest failed with rc: 0x%x", rc);
        return rc;
    }

    rc = Tss2_Sys_FlushContext (sys_ctx,
                                trialPolicySession->sessionHandle);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_FlushContext failed with rc: 0x%x", rc);
        return rc;
    }

    end_auth_session (trialPolicySession);
    return rc;
}

TSS2_RC
nv_rw_with_session (
    TSS2_SYS_CONTEXT *sys_ctx,
    const TPM2B_DIGEST *authPolicy,
    TPMA_NV nvAttributes,
    TPM2_SE session_type)
{
    TSS2_RC rc;
    TPM2B_AUTH  nvAuth = {
        .size = SECRET_SIZE,
        .buffer = SECRET_DATA,
    };
    SESSION *nvSession = NULL;
    TPM2B_NAME nvName;
    TPM2B_NONCE nonceCaller = { 0, };
    TPM2B_MAX_NV_BUFFER nvWriteData = {
        .size = NV_DATA_SIZE,
        .buffer = NV_DATA,
    };
    TPM2B_MAX_NV_BUFFER nvReadData = { .size = TPM2B_SIZE (nvReadData), };
    TPM2B_ENCRYPTED_SECRET encryptedSalt = { 0, };
    TPMT_SYM_DEF symmetric = {
        .algorithm = TPM2_ALG_NULL,
    };
    TSS2L_SYS_AUTH_RESPONSE nvRspAuths;
    TSS2L_SYS_AUTH_COMMAND nvCmdAuths = {
        .count = 1,
        .auths= {
            {
                .nonce = {
                    .size = 1,
                    .buffer = { 0xa5, },
                },
                .sessionHandle = TPM2_RH_PW,
                .sessionAttributes = TPMA_SESSION_CONTINUESESSION,
            }
        }
    };
    const TSS2L_SYS_AUTH_COMMAND auth_cmd_null_pwd = {
        .count = 1,
        .auths = {
            {
                .sessionHandle = TPM2_RH_PW,
            },
        },
    };

    rc = DefineNvIndex (sys_ctx,
                        TPM2_RH_PLATFORM,
                        &nvAuth,
                        authPolicy,
                        TPM20_INDEX_PASSWORD_TEST,
                        TPM2_ALG_SHA256,
                        nvAttributes,
                        32);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("DefineNvIndex failed with RC: 0x%x", rc);
        return rc;
    }

    /*
     * Add index and associated authorization value to
     * entity table.  This helps when we need
     * to calculate HMACs.
     */
    rc = AddEntity(TPM20_INDEX_PASSWORD_TEST, &nvAuth);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("AddEntity failed with RC: 0x%x", rc);
        return rc;
    }

    /* Get the name of the NV index. */
    rc = tpm_handle_to_name (sys_ctx,
                             TPM20_INDEX_PASSWORD_TEST,
                             &nvName);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("tpm_handle_to_name failed with RC: 0x%x", rc);
        return rc;
    }

    /*
     * Start HMAC or real (non-trial) policy authorization session:
     * it's an unbound and unsalted session, no symmetric
     * encryption algorithm, and SHA256 is the session's
     * hash algorithm.
     */
    rc = create_auth_session (&nvSession,
                              TPM2_RH_NULL,
                              0,
                              TPM2_RH_NULL,
                              0,
                              &nonceCaller,
                              &encryptedSalt,
                              session_type,
                              &symmetric,
                              TPM2_ALG_SHA256,
                              sys_ctx);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("create_auth_session failed with RC: 0x%x", rc);
        return rc;
    }

    /* set handle in command auth */
    nvCmdAuths.auths[0].sessionHandle = nvSession->sessionHandle;

    /*
     * Get the name of the session and save it in
     * the nvSession structure.
     */
    rc = tpm_handle_to_name (sys_ctx,
                             nvSession->sessionHandle,
                             &nvSession->name);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("tpm_handle_to_name failed with RC: 0x%x", rc);
        return rc;
    }

    /*
     * Now setup for writing the NV index.
     */
    if (session_type == TPM2_SE_POLICY) {
        rc = Tss2_Sys_PolicyAuthValue (sys_ctx,
                                       nvSession->sessionHandle,
                                       0,
                                       0);
        if (rc != TSS2_RC_SUCCESS) {
            LOG_ERROR ("Tss2_Sys_PolicyAuthValue failed with RC: 0x%x", rc);
            return rc;
        }
    }
    /* First call prepare in order to create cpBuffer. */
    rc = Tss2_Sys_NV_Write_Prepare (sys_ctx,
                                    TPM20_INDEX_PASSWORD_TEST,
                                    TPM20_INDEX_PASSWORD_TEST,
                                    &nvWriteData,
                                    0);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_NV_Write_Prepare failed with RC: 0x%x", rc);
        return rc;
    }

    /* Roll nonces for command */
    roll_nonces (nvSession, &nvCmdAuths.auths[0].nonce);

    /*
     * Complete command authorization area, by computing
     * HMAC and setting it in nvCmdAuths.
     */
    rc = compute_command_hmac(sys_ctx,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM2_RH_NULL,
                              &nvCmdAuths);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("compute_command_hmac failed with RC: 0x%x", rc);
        return rc;
    }

    /*
     * Finally!!  Write the data to the NV index.
     * If the command is successful, the command
     * HMAC was correct.
     */
    rc = TSS2_RETRY_EXP (Tss2_Sys_NV_Write (sys_ctx,
                            TPM20_INDEX_PASSWORD_TEST,
                            TPM20_INDEX_PASSWORD_TEST,
                            &nvCmdAuths,
                            &nvWriteData,
                            0,
                            &nvRspAuths));
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_NV_Write failed with RC: 0x%x", rc);
        return rc;
    }


    /* Roll nonces for response */
    roll_nonces (nvSession, &nvRspAuths.auths[0].nonce);

    /*
     * If the command was successful, check the
     * response HMAC to make sure that the
     * response was received correctly.
     */
    rc = check_response_hmac (sys_ctx,
                              &nvCmdAuths,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM2_RH_NULL,
                              &nvRspAuths);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("check_response_hmac failed with RC: 0x%x", rc);
        return rc;
    }

    if (session_type == TPM2_SE_POLICY) {
        rc = Tss2_Sys_PolicyAuthValue (sys_ctx,
                                       nvSession->sessionHandle,
                                       0,
                                       0);
        if (rc != TSS2_RC_SUCCESS) {
            LOG_ERROR ("Tss2_Sys_PolicyAuthValue failed with RC: 0x%x", rc);
            return rc;
        }
    }
    /* First call prepare in order to create cpBuffer. */
    rc = Tss2_Sys_NV_Read_Prepare (sys_ctx,
                                   TPM20_INDEX_PASSWORD_TEST,
                                   TPM20_INDEX_PASSWORD_TEST,
                                   NV_DATA_SIZE,
                                   0);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_NV_Read_Prepare failed with RC: 0x%x", rc);
        return rc;
    }

    roll_nonces (nvSession, &nvCmdAuths.auths[0].nonce);
    /* End the session after next command. */
    nvCmdAuths.auths[0].sessionAttributes &= ~TPMA_SESSION_CONTINUESESSION;

    /*
     * Complete command authorization area, by computing
     * HMAC and setting it in nvCmdAuths.
     */
    rc = compute_command_hmac (sys_ctx,
                               TPM20_INDEX_PASSWORD_TEST,
                               TPM20_INDEX_PASSWORD_TEST,
                               TPM2_RH_NULL,
                               &nvCmdAuths);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("compute_command_hmac failed with RC: 0x%x", rc);
        return rc;
    }

    /*
     * And now read the data back.
     * If the command is successful, the command
     * HMAC was correct.
     */
    rc = Tss2_Sys_NV_Read (sys_ctx,
                           TPM20_INDEX_PASSWORD_TEST,
                           TPM20_INDEX_PASSWORD_TEST,
                           &nvCmdAuths,
                           NV_DATA_SIZE,
                           0,
                           &nvReadData,
                           &nvRspAuths);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_NV_Read failed with RC: 0x%x", rc);
        return rc;
    }

    /* Roll nonces for response */
    roll_nonces (nvSession, &nvRspAuths.auths[0].nonce);

    /*
     * If the command was successful, check the
     * response HMAC to make sure that the
     * response was received correctly.
     */
    rc = check_response_hmac (sys_ctx,
                              &nvCmdAuths,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM20_INDEX_PASSWORD_TEST,
                              TPM2_RH_NULL,
                              &nvRspAuths);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("check_response_hmac failed with RC: 0x%x", rc);
        return rc;
    }

    /* Check that write and read data are equal. */
    if (memcmp ((void *)&nvReadData.buffer[0],
                (void *)&nvWriteData.buffer[0],
                nvReadData.size))
    {
        LOG_ERROR ("Data read not equal to data written.");
        return 1;
    }

    /*
     * Now cleanup:  undefine the NV index and delete
     * the NV index's entity table entry.
     */

    /* Undefine the NV index. */
    rc = Tss2_Sys_NV_UndefineSpace (sys_ctx,
                                    TPM2_RH_PLATFORM,
                                    TPM20_INDEX_PASSWORD_TEST,
                                    &auth_cmd_null_pwd,
                                    0);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR ("Tss2_Sys_NV_UndefineSpace failed with RC: 0x%x", rc);
        return rc;
    }

    /* Delete the NV index's entry in the entity table. */
    DeleteEntity (TPM20_INDEX_PASSWORD_TEST);

    /* Remove the real session from sessions table. */
    end_auth_session (nvSession);
    return rc;
}

TSS2_TCTI_CONTEXT *tcti_ctx = NULL;
TSS2_SYS_CONTEXT *sys_ctx = NULL;

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_ctx) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_ctx);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_ctx,
                         TSS2_SYS_CONTEXT **sys_ctx) {
    if (tcti_ctx == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_ctx = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_ctx == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_ctx, sys_context_size,
                                     tcti_ctx, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_ctx);
        *sys_ctx = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc =
        Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                               TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {

    if (tcti_ctx != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_ctx);
        free(tcti_ctx);
    }
    if (sys_ctx != NULL) {
        Tss2_Sys_Finalize(sys_ctx);
        free(sys_ctx);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Cleanup NV index if it exists
TSS2_RC cleanup_nv_index(TSS2_SYS_CONTEXT *sys_ctx) {
    const TSS2L_SYS_AUTH_COMMAND auth_cmd_null_pwd = {
        .count = 1,
        .auths = { { .sessionHandle = TPM2_RH_PW, }, },
    };
    
    TSS2_RC rc = Tss2_Sys_NV_UndefineSpace(sys_ctx,
                                           TPM2_RH_PLATFORM,
                                           TPM20_INDEX_PASSWORD_TEST,
                                           &auth_cmd_null_pwd,
                                           0);
    
    // Ignore errors since the index might not exist
    if (rc == TSS2_RC_SUCCESS) {
        LOG_INFO("Cleaned up existing NV index 0x%x", TPM20_INDEX_PASSWORD_TEST);
    }
    
    DeleteEntity(TPM20_INDEX_PASSWORD_TEST);
    return TSS2_RC_SUCCESS;
}

TSS2_RC
create_primary_key(TSS2_SYS_CONTEXT *sys_ctx, TPM2_HANDLE *primaryHandle)
{
    TSS2_RC rc;
    TPM2B_SENSITIVE_CREATE inSensitive = { 0, };
    TPM2B_PUBLIC inPublic = { 0, };
    TPM2B_DATA outsideInfo = { 0, };
    TPML_PCR_SELECTION creationPCR = { 0, };
    TPM2B_PUBLIC outPublic = { 0, };
    TPM2B_CREATION_DATA creationData = { 0, };
    TPM2B_DIGEST creationHash = { 0, };
    TPMT_TK_CREATION creationTk = { 0, };
    
    // Set up the primary key template for RSA 2048
    inPublic.publicArea.type = TPM2_ALG_RSA;
    inPublic.publicArea.nameAlg = TPM2_ALG_SHA256;
    inPublic.publicArea.objectAttributes = (TPMA_OBJECT_USERWITHAUTH |
                                            TPMA_OBJECT_RESTRICTED |
                                            TPMA_OBJECT_DECRYPT |
                                            TPMA_OBJECT_FIXEDTPM |
                                            TPMA_OBJECT_FIXEDPARENT |
                                            TPMA_OBJECT_SENSITIVEDATAORIGIN);
    
    inPublic.publicArea.parameters.rsaDetail.symmetric.algorithm = TPM2_ALG_AES;
    inPublic.publicArea.parameters.rsaDetail.symmetric.keyBits.aes = 128;
    inPublic.publicArea.parameters.rsaDetail.symmetric.mode.aes = TPM2_ALG_CFB;
    inPublic.publicArea.parameters.rsaDetail.scheme.scheme = TPM2_ALG_NULL;
    inPublic.publicArea.parameters.rsaDetail.keyBits = 2048;
    inPublic.publicArea.parameters.rsaDetail.exponent = 0;
    
    inPublic.publicArea.unique.rsa.size = 0;
    
    const TSS2L_SYS_AUTH_COMMAND auth_cmd_null_pwd = {
        .count = 1,
        .auths = {
            {
                .sessionHandle = TPM2_RH_PW,
            },
        },
    };
    
    rc = Tss2_Sys_CreatePrimary(sys_ctx,
                                TPM2_RH_OWNER,
                                &auth_cmd_null_pwd,
                                &inSensitive,
                                &inPublic,
                                &outsideInfo,
                                &creationPCR,
                                primaryHandle,
                                &outPublic,
                                &creationData,
                                &creationHash,
                                &creationTk,
                                NULL,
                                NULL);
    
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR("Tss2_Sys_CreatePrimary failed with RC: 0x%x", rc);
    } else {
        LOG_INFO("Primary key created successfully with handle: 0x%x", *primaryHandle);
    }
    
    return rc;
}

int main(int argc, char *argv[])
{
    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_ctx);
    CHECK_RC(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_ctx, &sys_ctx);
    CHECK_RC(rc);

    // Startup TPM
    rc = startup_tpm(sys_ctx);
    CHECK_RC(rc);
    
    TPM2B_DIGEST authPolicy = { 0, };
    TPMA_NV nvAttributes;
    
    // Cleanup any existing NV index
    cleanup_nv_index(sys_ctx);
    
    LOG_INFO ("=== Basic HMAC session test ===");
    printf("Starting basic HMAC session test...\n");
    nvAttributes = TPMA_NV_AUTHREAD | TPMA_NV_AUTHWRITE | TPMA_NV_PLATFORMCREATE;
    rc = nv_rw_with_session (sys_ctx, &authPolicy, nvAttributes, TPM2_SE_HMAC);
    if (rc != TSS2_RC_SUCCESS) {
        LOG_ERROR("Basic HMAC session test failed");
        cleanup(rc);
    }
    LOG_INFO ("Basic HMAC session test completed successfully");
    printf("Basic HMAC session test completed successfully!\n");
    
    // Cleanup before next test
    cleanup_nv_index(sys_ctx);

    cleanup(TSS2_RC_SUCCESS);
    return 0;
}
